library(dplyr)
library(tidyr)
library(ggplot2)
library(viridis)
library(ggpubr)
library(cowplot)
library(gridExtra)
library(stringr)

# Dichotomize signatures
# Select signatures present in >n% of samples
# Dichotomize the signatures (Presence (signatures>0)="1",absence (signatures==0)="0")
# If signature is present in more than 75% cases, dichotomize by the median
# Return data frame with metadata and dichotomized signatures

# metadata: Data frame with sample metadata
# sigsn: Data frame with signature attributions
# sigcutoff: cutoff in percentage. Default is 0% (i.e. do not exclude any signature)

dichotomizeSigs <- function(metadata, sigsn, sigcutoff = 0) {
  for (i in names(sigsn)[-1]) {
    sigsn[i] <- ifelse(metadata[metadata$donor_id %in% sigsn$donor_id, i] > 0, 1, 0)
  }
  # Select signatures above cutoff frequency
  sigsp <- colSums(sigsn[, -1]) / nrow(sigsn) * 100
  signatures <- names(sigsp[sigsp > sigcutoff])
  # Dichotomize selected signatures
  for (i in signatures) {
    if (length(which(metadata[, i] > 0)) <= nrow(sigsn) * 0.75) {
      metadata[paste0(i, "_cat")] <- as.factor(ifelse(metadata[, i] > 0, "1", "0"))
      if (0 %in% levels(metadata[, paste0(i, "_cat")])) {
        metadata[, paste0(i, "_cat")] <- relevel(metadata[, paste0(i, "_cat")], ref = "0")
      } else {
        print(paste("all samples are positive for", i))
      }
    } else {
      print(paste("categorizing", i, "by sample median"))
      metadata[paste0(i, "_cat")] <- as.factor(ifelse(metadata[, i] >= median(metadata[which(metadata[, i] > 0), i]), "1", "0"))
      if (0 %in% levels(metadata[, paste0(i, "_cat")])) {
        metadata[, paste0(i, "_cat")] <- relevel(metadata[, paste0(i, "_cat")], ref = "0")
      } else {
        print(paste("all samples are positive for", i))
      }
    }
  }
  return(metadata)
}


# Perform Kruskal-Wallis Test and Generate Plots
# This function conducts Kruskal-Wallis tests between variables of interest and signature burdens,
# generates boxplots with signature frequencies, and returns the test results and plots.

# dat: Data frame containing variables  and signatures of interest.
# signatures: Character vector specifying signature names.
# factors: Character vector specifying variables of interest.
# pval_cutoff: p-value threshold for plotting a specific test.
# output.kw: Data frame where Kruskal-Wallis test results will be stored. Ddefault is an empty data frame.
# cut_y: cut y axis at 1.25 × upper whisker. Default is TRUE.
# y_labs: Label for the y-axis in the plots. Default is "Mutation burden".
# x_labs: Label for the x-axis in the plots. Default is the same as the factor variable.

plot_kw <- function(dat, signatures, factors, pval_cutoff = 0.05, output.kw = data.frame(),
                    cut_y = TRUE, y_labs = "Mutation burden", x_labs = NULL) {
  plots <- list()

  # Loop over each signature and factor
  for (s in signatures) {
    for (f in factors) {
       x_labs_plot <- if (!is.character(x_labs)) f else x_labs

      # Check if the factor has less than 2 categories, if so, print a warning and skip
      if (dat[, f] %>% unique() %>% length() < 2) {
        print(paste("Warning: Variable", f, "has been removed due to <2 categories"))
      } else {
        kw <- tryCatch(
          {
            kruskal.test(dat[, s], dat[, f])
          },
          error = function(e) {
            NULL
          }
        )

        # If there was an error in computing the test, move to the next factor
        if (is.null(kw)) {
          print(paste("Variable", f, "has been removed due to ERROR in KW test"))
          next
        }

        result.var <- data.frame(
          signature = paste(s),
          factor = paste(f),
          pval = kw$p.value
        )

        output.kw <- rbind(output.kw, result.var)

        # If p-value is less than the cutoff, proceed with plotting
        if (kw$p.value < pval_cutoff) {
          # Remove NAs
          subset <- dat[!is.na(dat[f]) & !is.na(dat[s]), ]

          # Prepare table with frequencies
          propdata <- dat %>%
            group_by(get(f), get(paste0(s, "_cat"))) %>%
            tally() %>%
            na.omit()
          names(propdata) <- c(f, paste0(s, "_cat"), "n")
          propdata <- spread(propdata, paste0(s, "_cat"), "n")
          propdata[is.na(propdata)] <- 0

          # Of all samples are positive for signature i.e. there is no 0 value for sigs, print 100% positivity in plot
          if ("0" %in% colnames(propdata)) {
            propdata <- propdata %>% mutate(prop = paste0(`1`, "/", `0` + `1`), perc = `1` / (`0` + `1`) * 100)
          } else {
            propdata <- propdata %>% mutate(prop = paste0(`1`, "/", `1`), perc = `1` / (`1`) * 100)
          }

          # Calculate maximum y-axis limit
          ylim_max <- subset %>%
            group_by(get(f)) %>%
            summarise(upper_whisker = boxplot.stats(get(s))$stats[c(5)]) %>%
            pull(upper_whisker) %>%
            max()
          if (cut_y == F | ylim_max == 0) {
            ylim_max <- max(subset[, s])
          }

          # plot
          p_top <- ggplot(data = subset, aes(x = get(f), y = get(s), fill = get(f))) +
            geom_violin(aes(color = get(f)), width = 1, alpha = .8, linewidth = .3) +
            geom_boxplot(color = "#585858", alpha = 0.15, linewidth = .3, staplewidth = .2, outlier.shape = NA) +
            geom_jitter(fill = "#585858", size = .8, stroke = 0, alpha = 0.25, width = .1, height = .1) +
            labs(
              title = s,
              subtitle = ifelse(nrow(propdata) == 2, paste("Wilcoxon, p =", signif(kw$p.value, 3)), paste("Kruskal-Wallis, p =", signif(kw$p.value, 3))),
              y = y_labs
            ) +
            coord_cartesian(ylim = c(0, ylim_max * 1.25)) +
            scale_fill_viridis(discrete = TRUE, end = .8) +
            scale_color_viridis(discrete = TRUE, end = .8) +
            theme_minimal() +
            theme(
              legend.position = "none",
              plot.title = element_text(face = "bold", size = 12, hjust = 0.5),
              plot.subtitle = element_text(size = 10, hjust = 0.5),
              axis.title.x = element_blank(),
              axis.text.x = element_blank(),
              axis.text = element_text(size = 12),
              axis.line.y = element_line(size = .3,color = "#404040"),
              panel.grid = element_blank(),
              plot.margin = unit(c(0.5, 0, 0.1, 0), "lines")
            )

          p_bottom <- ggplot(propdata, aes(x = get(f), y = perc)) +
            geom_col(width = 0.5, alpha = 0.8) +
            geom_text(aes(y = perc + 20, x = get(f), label = prop), size = 3) +
            scale_y_continuous(limits = c(0, 100), breaks = c(0, 50, 100)) +
            labs(
              title = s,
              x = str_to_title(str_replace_all(x_labs_plot, "_", " ")),
              y = "Frequency"
            ) +
            theme_minimal() +
            theme(
              legend.position = "none",
              plot.title = element_blank(),
              axis.text = element_text(size = 12),
              axis.text.x = element_text(angle = 45, hjust = 1),
              axis.line = element_line(size = .3,color = "#404040"),
              axis.ticks = element_line(size = .3,color = "#404040"),
              panel.grid = element_blank(),
              plot.margin = unit(c(0.5, 0.5, 0.5, 0.5), "lines")
            )

          # Combine top and bottom panels and store
          p_all <- ggarrange(p_top, p_bottom,
            heights = c(0.55, 0.45),
            ncol = 1, align = "v"
          )
          print(p_all)
          plots[[paste(s, f, sep = "_")]] <- p_all
        }
      }
    }
  }

  # Calculate adjusted p-values
  output.kw <- output.kw %>% mutate(p_adj = ifelse(pval * length(signatures) < 1, pval * length(signatures), 1))
  return(list("output.kw" = output.kw, "plots" = plots))
}

# Import CN data as data frame
# id: Sample ID
# DIR: Directory where battenberg files are stored

battenberg_import <- function(id, dir) {
  input_files <- str_subset(list.files(dir), id)
  cn_name <- str_subset(input_files, "battenberg.summary")
  cn <- read.csv(str_c(dir, cn_name), header = F, col.names = c("X", "chromosome", "start", "end", "totalCNnorm", "minCNnorm", "totalCNtum", "minCNtum")) %>%
    filter(end - start > 1000000) %>%
    dplyr::select(-X)
}

# Import CN data as GRanges object
# id: Sample ID
# DIR: Directory where battenberg files are stored

GRanges_import <- function(id, dir) {
  input_files <- str_subset(list.files(dir), id)
  cn_name <- str_subset(input_files, "battenberg.summary")
  cn_battenberg <- read.csv(str_c(dir, cn_name), header = F) %>%
    mutate(strand = "*") %>%
    filter(V4 - V3 > 1000000)
  bb <- GRanges(
    seqnames = as(cn_battenberg$V2, "Rle"), ranges = IRanges(start = cn_battenberg$V3, end = cn_battenberg$V4), strand = as(cn_battenberg$strand, "Rle"),
    total_cn = cn_battenberg$V7, minor_cn = cn_battenberg$V8
  )
  return(bb)
}
